# Loss functions

import torch
import torch.nn as nn

from utils.general import bbox_iou
from utils.torch_utils import is_parallel


def smooth_BCE(eps=0.1):  # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
    # return positive, negative label smoothing BCE targets
    return 1.0 - 0.5 * eps, 0.5 * eps


class BCEBlurWithLogitsLoss(nn.Module):
    # BCEwithLogitLoss() with reduced missing label effects.
    def __init__(self, alpha=0.05):
        super(BCEBlurWithLogitsLoss, self).__init__()
        self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none')  # must be nn.BCEWithLogitsLoss()
        self.alpha = alpha

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)
        pred = torch.sigmoid(pred)  # prob from logits
        dx = pred - true  # reduce only missing label effects
        # dx = (pred - true).abs()  # reduce missing label and false label effects
        alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
        loss *= alpha_factor
        return loss.mean()


class FocalLoss(nn.Module):
    # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
        super(FocalLoss, self).__init__()
        self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = loss_fcn.reduction
        self.loss_fcn.reduction = 'none'  # required to apply FL to each element

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)
        # p_t = torch.exp(-loss)
        # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power for gradient stability

        # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
        pred_prob = torch.sigmoid(pred)  # prob from logits
        p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
        modulating_factor = (1.0 - p_t) ** self.gamma
        loss *= alpha_factor * modulating_factor

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:  # 'none'
            return loss


class QFocalLoss(nn.Module):
    # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
        super(QFocalLoss, self).__init__()
        self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = loss_fcn.reduction
        self.loss_fcn.reduction = 'none'  # required to apply FL to each element

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)

        pred_prob = torch.sigmoid(pred)  # prob from logits
        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
        modulating_factor = torch.abs(true - pred_prob) ** self.gamma
        loss *= alpha_factor * modulating_factor

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:  # 'none'
            return loss


class ComputeLoss:
    """
        其实还是通过预测图（detect出来的三个预测图，分别【10（一个batch图像10个）,3(anchor数量),80（gridx）,80(gridy),8(xywh,置信度，类别）】）
        通过self.build_target来预测哪些格子和对应的anchor要被选择（变量为indices，包含格子信息），将图片中每个格子的值提取出来，
        然后经过回归函数pxy = ps[:, :2].sigmoid() * 2. - 0.5 #【1672，2】来进行回归

        这样就可以得到xywh的值，求得与tbox相关的损失值
    """
    # Compute losses
    def __init__(self, model, autobalance=False):
        super(ComputeLoss, self).__init__()
        device = next(model.parameters()).device  # get model device
        #获得超参数
        h = model.hyp  # hyperparameters

        # Define criteria
        #定义类别和目标性损失函数！！
        """
            cls: 0.5  # cls loss gain # 分类损失的系数
            cls_pw: 1.0  # cls BCELoss positive_weight #分类BCELoss中正样本的权重
            obj: 1.0  # obj loss gain (scale with pixels) # 有无物体损失的系数
            obj_pw: 1.0  # obj BCELoss positive_weight # 有无物体BCELoss中正样本的权重
            
            https://blog.csdn.net/qq_22210253/article/details/85222093
            BCEWithLogitsLoss就是把Sigmoid-BCELoss合成一步。
            pos_weight是正样本的权重，解决样本不均衡的问题
        """
        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))

        # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
        #标签平滑
        self.cp, self.cn = smooth_BCE(eps=0.0)

        # Focal loss  focal loss gamma   默认为0没有使用
        g = h['fl_gamma']  # focal loss gamma
        if g > 0:
            BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)

        det = model.module.model[-1] if is_parallel(model) else model.model[-1]  # Detect() module
        #设置三个特征图对应的损失系数!!!!!!
        self.balance = {3: [4.0, 1.0, 0.4], 4: [4.0, 1.0, 0.25, 0.06], 5: [4.0, 1.0, 0.25, 0.06, .02]}[det.nl]#和detect的层数有关如果detect不止一层的话
        self.ssi = (det.stride == 16).nonzero(as_tuple=False).item()  # stride 16 index
        self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance
        for k in 'na', 'nc', 'nl', 'anchors':
            setattr(self, k, getattr(det, k))

    #像函数一样调用
    def __call__(self, p, targets):  # predictions, targets, model
        device = targets.device
        #初始化各部分损失  lcls:类别损失    lbox:anchorbox损失   lobj：置信度损失
        lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
        """
            #在每个yolo层将预设的anchor和ground truth进行匹配，得到正样本anchor
            :return: tcls表示类别;
            tbox表示标记的box和生成的box的左边（x,y,w,h)
            indices:图像索引，选取的anchor的索引（每层特征图中每个网格初始化3个），每层特征图上选取的网格点坐标（i，j）
            anch选取的anchor的索引
        """
        tcls, tbox, indices, anchors = self.build_targets(p, targets)  # targets #排除一些anchors，并调整格式方便后面计算loss

        # Losses 预测每个输出
        for i, pi in enumerate(p):  # layer index, layer predictions 【pi第一个.shape】10（一个batch图像10个）,3(anchor数量),80（gridx）,80(gridy),8(xywh,置信度，类别）
            b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx index值    取出对应detect层的正样本 都是向量【1672】
            tobj = torch.zeros_like(pi[..., 0], device=device)  # target obj 初始化置信度！！！！！！！！！！！！！！！！ 【10，3，80，80】

            n = b.shape[0]  # number of targets
            if n:
                ps = pi[b, a, gj, gi]  # prediction subset corresponding to targets 提取该层某格子的信息 1672组向量选择 最后结果为[1672,8]

                # Regression
                # 这个计算的是预测框的xywh 这个就是目标框回归公式   txty=pxy 相对于左上角格子的偏移值
                pxy = ps[:, :2].sigmoid() * 2. - 0.5 #【1672，2】
                pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
                pbox = torch.cat((pxy, pwh), 1)  # predicted box【1672，8】预测框
                ## 下面就是构造损失函数让预测结果通过怎样的优化去更接近真值了！！！！！！！！！！！！！！！！！！！！！！！！！！
                #计算边框损失，计算的是ciou损失
                iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)  # iou(prediction, target) #求iou  返回[1672]和box数一样
                lbox += (1.0 - iou).mean()  # iou loss 1-ciou的平均值 这就是CIOU的损失值

                # Objectness
                #接下来就要对这12个行人的置信度进行设置成1（有权重self.gr * iou.detach()），因为前面我们是初始化成全0。
                #故将预测框和bbox的iou作为权重乘到conf分支，用于正样本的置信度损失赋值
                """
                    因为只有12个行人对应的预测框的理论置信度被设置成1，其他的预测框的理论置信度都是0。在实际计算过程中，所有YOLO层的实际预测框都会去求置信度损失，但在这些预测框里面其实只有少数是真的含有目标的，
                    这就会导致没有目标的负样本在这个损失部分里占了主导地位，容易使系统的训练方向跑偏，系统变得不怎么敢去预测行人这个结果，
                    所以才会有focal loss损失的出现解决这个问题，但是加入focal loss之后系统又变得过于大胆预测行人，会出现更多误检。
                """
                tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype)  # iou ratio

                # Classification 判断是不是多分类
                if self.nc > 1:  # cls loss (only if multiple classes)
                    t = torch.full_like(ps[:, 5:], self.cn, device=device)  # targets
                    t[range(n), tcls[i]] = self.cp
                    lcls += self.BCEcls(ps[:, 5:], t)  # BCE 损失函数计算

                # Append targets to text file
                # with open('targets.txt', 'a') as file:
                #     [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]

            obji = self.BCEobj(pi[..., 4], tobj)#损失函数计算
            lobj += obji * self.balance[i]  # obj loss 每个特征图都会预设不同的置信度损失系数 为了平衡各个特征图的差异
            if self.autobalance:
                self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()

        if self.autobalance:
            self.balance = [x / self.balance[self.ssi] for x in self.balance]
        lbox *= self.hyp['box']#平衡各个损失之间的差异
        lobj *= self.hyp['obj']
        lcls *= self.hyp['cls']
        bs = tobj.shape[0]  # batch size

        loss = lbox + lobj + lcls #总损失
        return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()

    def build_targets(self, p, targets):
        """https://blog.csdn.net/a1874738854/article/details/112789533
            简单来说就去去除部分不合规则的anchor减少计算量，并且整理剩余的anchors，方便计算loss

            build_targets函数用于获得在训练时计算loss函数所需要的目标框，即被认为是正样本，
            与yolov3/v4不同，yolov5支持跨网格预测！！！
            对于任何一个bbox，三个输出预测特征层都可能有先验框anchors进行匹配
            该函数输出的正样本比传入的target(gt框）数目多
            具体处理过程：
            （1）对于任何一层计算当前bbox和但浅层anchor的匹配程度，不采用iou，而是shape比例；
            如果anchor和bbox的宽高比差距大于4则认为不匹配，此时忽略相应的bbox，即当做是背景；
            （2）然后对bbox计算落在的网格所有anchors都计算loss（并不是直接和GT框比较计算loss）
            注意此时落在网格不再是一个，而是附近的多个，这样就增加了正样本数，可能存在有些bbox在三个尺度都预测的情况
            另外，yolov5也没有conf分支忽略阈值（ignore_thresh）的操作，而yolov3/v4有
            target====="tensor([[0.00000, 2.00000, 0.49836, 0.11987, 0.41351, 0.23974],
                [0.00000, 2.00000, 0.92385, 0.17014, 0.15230, 0.34028],
                [0.00000, 0.00000, 0.09377, 0.49597, 0.01495, 0.01941],
                ...,
                [9.00000, 0.00000, 0.10651, 0.85327, 0.02445, 0.03094],
                [9.00000, 0.00000, 0.05311, 0.84778, 0.03543, 0.04192],
                [9.00000, 0.00000, 0.03889, 0.83281, 0.03194, 0.03393]], device='cuda:0')
        :param p:models输出的预测的特征图的大小
        :param targets:标注的label值,targets最后的7表示【batchindex,类型值，x,y,w,h,最后一个是anchor?】 【3，312，7】
        :return: tcls表示类别;
            tbox表示标记的box和生成的box的左边（x,y,w,h)
            indices:图像索引，选取的anchor的索引（每层特征图中每个网格初始化3个），每层特征图上选取的网格点坐标（i，j）
            anch选取的anchor的索引
        """
        # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
        na, nt = self.na, targets.shape[0]  # number of anchors, targets 是一个batch总的数目
        """
            初始化每个batch box的信息：
            tcls表示类别;
            tbox表示标记的box和生成的box的左边（x,y,w,h)
            indices:图像索引，选取的anchor的索引（每层特征图中每个网格初始化3个），每层特征图上选取的网格点坐标（i，j）
            anch选取的anchor的索引
        """
        tcls, tbox, indices, anch = [], [], [], [] #最终需要返回的
        gain = torch.ones(7, device=targets.device)  # normalized to gridspace gain 初始化为1
        #ai.shape=(na,nt)生成anchor索引
        #anchor索引和后面有用，用于表示当前bbox和当前层的哪个anchor进行匹配
        ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
        #先repeat target和当前层的anchor个数一样，相当于bbox编程了三个，然后和3个anchors相匹配 也就是n*7变为3*n*7
        targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)  # append anchor indices

        #设置网格中心的偏移量
        g = 0.5  # bias
        #附近的网格 【1，0】x增加1的网格，也就是当前网格上面一个网格
        off = torch.tensor([[0, 0],
                            [1, 0], [0, 1], [-1, 0], [0, -1],  # j,k,l,m
                            # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
                            ], device=targets.device).float() * g  # offsets
        #对每个检测层进行处理 三个预测shape【0:[10,3,80,80,8];;;1:[10,3,40,40,8];;;2:[10,3,20,20,8]】  8应该就是5+3坐标 3应该是因为图片一般都是rbg，这个10确实不清楚
        for i in range(self.nl):#三个尺度的特征预测图的分支
            anchors = self.anchors[i]#当前分支的anchor大小，在yaml文件中预置的大小（已经除以了当前特征图所对应的stride，stride也就是8,16,32降采样率）[14,27]/8=【1.75，3.375】
            gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  #p是网络输出， xyxy gain 这句是为提取特征图的大小

            # Match targets to anchors 标注的label值
            t = targets * gain#将标签的xywh从基于0~1映射到特征图，targets的xywh本身也就是归一化的尺度，需要变成特征图尺寸，之后的操作会不断削减t的内容
            #对每个输出层都单独匹配
            if nt:
                # 第一次过滤预测的的bbox  下面两行代码就是采用shape匹配规则（非iou）也就是计算宽高比和阈值比较
                #bw=(sigmoid(tw))^2 bw范围在0-4
                r = t[:, :, 4:6] / anchors[:, None]  # wh ratio
                j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t']  # compare 为4
                # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                t = t[j]  # filter 筛选<self.hyp['anchor_t']预测的bbox框 过滤规则是没有考虑xy的

                # Offsets
                """
                    实际并不会用4+1个框，而会把相对于各个网格的左上角x<0.5,y<0.5和相对于右下角的x<0.5,y<0.5的框提取出来（四舍五入，与label最近的），也就是j,k,l,m（上下左右？）（在off变量中有说明）
                    在选取gij（也就是标签框分配给的网格）的时候对这个部分的框都做一个偏移（剪去offset），这也就是下面的gij的操作=(gxy-offset).long()的操作
                    再将这四个部分的框和原始的gxy拼接在一起，总共就是五个部分
                    yolov3/4仅仅使用当前的网格进行回归，yolov5相较于v4限制了范围-0.5~1.5
                    xy.sigmoid()*2-0.5+cx 
                    
                    对于剩下的bbox，计算其落在哪个网格内，同时利用四舍五入规则，找出最近的两个网格，
                    将这三个网格都认为是负责预测该bbox的，可以发现粗略估计正样本数相比前yolo系列，增加了三倍。code如下：
                    剩下的t（前面通过shape过滤后），现在用offset进行过滤，只找到相关的
                """
                #左上角为轴中心
                gxy = t[:, 2:4]  # grid xy 预测的中心点坐标  t的shape为[558,7]  取出的就是2,3是x,y的坐标
                #得到右下角为轴中心的坐标值
                gxi = gain[[2, 3]] - gxy  # inverse
                # 上下左右 分别判断box的（x，y）坐标是否大于1，并距离网格左上角的距离（准确的说是y距离网格上边或x距离网格左边的距离）距离小于0.5，如果（x，y）中满足上述两个条件，则选中
                j, k = ((gxy % 1. < g) & (gxy > 1.)).T #这里用到了浮点数取摸，对于两个浮点数a,b a%b=a-n*b 其中n为不超过a/b的最大整数
                l, m = ((gxi % 1. < g) & (gxi > 1.)).T
                j = torch.stack((torch.ones_like(j), j, k, l, m))#挑选一些框出来，五个558 拼出来shape[5,558]  因为预设的就是5个（上下左右中）
                t = t.repeat((5, 1, 1))[j]#然后重复5倍 【558，7】->[5,558,7]  用j来过滤 shape变成[5*过滤剩余,7]  找出左上右下最近的偏移
                offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j] #选择最近的三个！！！！ [1,558,2]+[5,1,2]=【5，558，2】然后j筛选变成和t一样的shape  把所有为1的轴扩充到与两个轴大小相同然后相加
            else:
                t = targets[0]
                offsets = 0

            # Define
            """
                对每个bbox找出对应的正样本anchor，
                其中包括b表示当前bbox属于batch内部的第几张图片，
                a表示当前bbox和当前层的第几个anchor匹配上，
                gi,gj是对应负责预测该bbox的网格坐标
                gxy是不考虑offset或者说v3/v4里面设定的该bbox的负责预测网格中心点坐标xy
                gwh是对应的bbox wh
                c是该bbox类别
         """
            b, c = t[:, :2].long().T  # image, class
            #中心点回归的坐标
            gxy = t[:, 2:4]  # grid xy
            #宽高回归
            gwh = t[:, 4:6]  # grid wh
            gij = (gxy - offsets).long()#计算当前label落在哪个网格上！！！！！！！！！！！！！！！！！！！！
            gi, gj = gij.T  # grid xy indices 网格的索引值gi是x的,gj是y的索引值，158个框

            # Append
            #a为anchor的索引
            a = t[:, 6].long()  # anchor indices
            #indices:图像索引，选取的anchor的索引（每层特征图中每个网格初始化3个），每层特征图上选取的网格点坐标（i，j）     【3，4，1133】
            indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
            # tbox表示标记的box（x,y,w,h) 【3，1133，4】
            tbox.append(torch.cat((gxy - gij, gwh), 1))  # box 作为预测框的坐标值
            #anch选取的anchor的索引 【3，1133，2】
            anch.append(anchors[a])  # anchors
            tcls.append(c)  # class 【3，1133，1】
        """
        由 indices 和 anch 可以还原成匹配的 Anchor 所在的位置，而由 indices 和 tbox 可以还原出匹配的 GT 所在的位置！！！！！！！！！！！！！！！！！！！！！！！！！！！
        """
        return tcls, tbox, indices, anch
